"""
https://blog.csdn.net/menghaocheng/article/details/102783705

【TF2.0-CNN】迁移学习（将inceptionV3应用到猫狗分类）
"""

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf

# local_weights_file = '/tmp/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5'
local_weights_file = '../../../../../large_data/model/inceptionV3/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5'

pre_trained_model = tf.keras.applications.inception_v3.InceptionV3(input_shape=(150, 150, 3),
                                                                   include_top=False,
                                                                   weights=None)

pre_trained_model.load_weights(local_weights_file)
for layer in pre_trained_model.layers:
    layer.trainable = False

last_layer = pre_trained_model.get_layer('mixed7')
print('last layer output shape: ', last_layer.output_shape)
last_output = last_layer.output
x = tf.keras.layers.Flatten()(last_output)
x = tf.keras.layers.Dense(1024, activation='relu')(x)
x = tf.keras.layers.Dropout(0.2)(x)
x = tf.keras.layers.Dense(1, activation='sigmoid')(x)

model = tf.keras.Model(pre_trained_model.input, x)

model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=0.0001),
              loss='binary_crossentropy',
              metrics=['acc'])

train_datagen = ImageDataGenerator(rescale=1. / 255.,
                                   rotation_range=40,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True)

dir = "../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/train"
train_generator = train_datagen.flow_from_directory(dir,
                                                    batch_size=20,
                                                    class_mode='binary',
                                                    target_size=(150, 150))

test_datagen = ImageDataGenerator(rescale=1.0 / 255.)
dir = "../../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/validation"
validation_generator = test_datagen.flow_from_directory(dir,
                                                        batch_size=20,
                                                        class_mode='binary',
                                                        target_size=(150, 150))

history = model.fit_generator(
    train_generator,
    validation_data=validation_generator,
    steps_per_epoch=100,
    epochs=2,
    validation_steps=50,
    verbose=1)
